"""
难度：中等
给定一个由 0 和 1 组成的矩阵 mat ，请输出一个大小相同的矩阵，其中每一个格子是 mat 中对应位置元素到最近的 0 的距离。
两个相邻元素间的距离为 1 。
示例 1：
输入：mat = [[0,0,0],[0,1,0],[0,0,0]]
输出：[[0,0,0],[0,1,0],[0,0,0]]
示例 2：
输入：mat = [[0,0,0],[0,1,0],[1,1,1]]
输出：[[0,0,0],[0,1,0],[1,2,1]]
提示：
m == mat.length
n == mat[i].length
1 <= m, n <= 104
1 <= m * n <= 104
mat[i][j] is either 0 or 1.
mat 中至少有一个 0 """
import collections

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        row_count = len(mat)
        col_count = len(mat[0])
        dist_map = [[0 for _ in range(col_count)] for _ in range(row_count)]
        zeroes_pos = []
        for i in range(row_count):
            for j in range(col_count):
                if mat[i][j] == 0:
                    zeroes_pos.append((i, j))

        directions = {(1, 0), (-1, 0), (0, 1), (0, -1)}
        queue = collections.deque(zeroes_pos)
        visited = set(zeroes_pos)

        while queue:
            i, j = queue.popleft()
            for direction in directions:
                new_i = i + direction[0]
                new_j = j + direction[1]
                if 0 <= new_i < row_count and 0 <= new_j < col_count and (new_i, new_j) not in visited:
                    dist_map[new_i][new_j] = dist_map[i][j] + 1
                    queue.append((new_i, new_j))
                    visited.add((new_i, new_j))
        return dist_map
